from fastapi import FastAPI, File, UploadFile
from typing import Dict
import uvicorn
import fitz  # PyMuPDF
import os
import random
import string
import datetime
from fastapi.responses import FileResponse

# 创建FastAPI应用实例
app = FastAPI()

def classify_pdf_pymupdf(pdf_path):
    """
    使用 PyMuPDF 判断 PDF 类型：文本型、图像型、可搜索 (OCR) 或空。

    Args:
        pdf_path (str): PDF 文件路径。

    Returns:
        str: "Text-based", "Image-based", "Searchable (OCR)", "Empty", or "Error: [message]"
    """
    if not os.path.exists(pdf_path):
        return f"Error: File not found at {pdf_path}"
    if os.path.getsize(pdf_path) == 0:
        return "Empty" # Handle zero-byte files

    doc = None  # Initialize doc to None
    try:
        doc = fitz.open(pdf_path)
        
    except Exception as e:
        # Catch specific fitz exceptions if needed, e.g., fitz.fitz.FileNotFoundError is redundant with os.path.exists
        # Catching generic Exception might hide issues, but okay for broad example.
        # A common one is needing a password:
        if "password" in str(e).lower():
             return "Error: PDF is password protected"
        return f"Error: Could not open PDF - {e}"

    if doc.needs_pass:
        doc.close()
        return "Error: PDF is password protected"

    if doc.page_count == 0:
        doc.close()
        return "Empty"

    total_text_length = 0
    text_page_count = 0  # Pages with significant text
    image_page_count = 0 # Pages that look like images AND lack significant text
    ocr_candidate_page_count = 0 # Pages that look like images BUT have significant text

    # --- Heuristics (可调整的阈值) ---
    min_text_per_page_threshold = 30  # 页面至少有多少字符才算有意义的文本
    image_area_ratio_threshold = 0.85 # 页面图像总面积占页面面积多大比例才算"像图像"
    # --- End Heuristics ---

    try:
        for page_num in range(doc.page_count):
            page = doc.load_page(page_num)
            page_area = page.rect.width * page.rect.height
            if page_area == 0: continue # Skip pages with no area

            # 1. Extract text
            page_text = page.get_text("text").strip()
            page_text_len = len(page_text)
            total_text_length += page_text_len
            has_significant_text = page_text_len >= min_text_per_page_threshold

            # 2. Check image coverage
            is_mostly_image = False
            try:
                img_list = page.get_images(full=True)
                total_image_area = 0
                for img_info in img_list:
                    xref = img_info[0]
                    try:
                        base_image = doc.extract_image(xref)
                        if base_image and base_image.get("width") and base_image.get("height"):
                            total_image_area += base_image["width"] * base_image["height"]
                    except Exception:
                        pass # Ignore errors extracting single images
                if page_area > 0 : # Avoid division by zero for empty pages
                    image_ratio = total_image_area / page_area
                    is_mostly_image = image_ratio >= image_area_ratio_threshold
            except Exception as img_ex:
                print(f"Warning: Could not process images on page {page_num}: {img_ex}")


            # 3. Classify page based on heuristics
            if has_significant_text:
                text_page_count += 1
                if is_mostly_image:
                    # Looks like an image, but has text -> Strong OCR candidate
                    ocr_candidate_page_count += 1
            elif is_mostly_image:
                # Looks like an image and has no significant text -> Image page
                image_page_count += 1
            else:
                # Doesn't look like an image, and has no significant text
                # Could be a blank page, or vector graphics without much text
                # We won't strongly classify it as image or text based on this alone
                pass


        # --- Final Classification Logic (可调整) ---
        num_pages = doc.page_count
        doc.close() # Close the document once done processing pages

        # Prioritize OCR candidate count
        if ocr_candidate_page_count >= num_pages * 0.5: # If at least half the pages look like OCR'd images
            return {"type": "Searchable (OCR)", "pages": num_pages}

        # If not predominantly OCR, check text vs image pages
        # Consider a PDF text-based if a significant portion of pages have text
        if text_page_count >= num_pages * 0.3: # If at least 30% of pages have significant text
             # Double check if it might still be OCR despite not meeting the first OCR rule
             if image_page_count >= num_pages * 0.5 and ocr_candidate_page_count > 0:
                  return {"type": "Searchable (OCR)", "pages": num_pages} # Looks like images overall, but some text found
             else:
                  return {"type": "Text-based", "pages": num_pages}

        # Consider image-based if very few text pages AND many image-like pages
        if image_page_count >= num_pages * 0.6 and text_page_count < num_pages * 0.2:
            return {"type": "Image-based", "pages": num_pages}

        # Ambiguous cases - rely on total text as a fallback, could be mixed content
        if total_text_length > min_text_per_page_threshold * num_pages * 0.1: # Check if total text is somewhat significant
            # If overall text exists, but pages looked like images, lean towards OCR
            if image_page_count + ocr_candidate_page_count >= num_pages * 0.7:
                 return {"type": "Searchable (OCR)", "pages": num_pages}
            else: # Otherwise, assume it's text-based despite ambiguity
                 return {"type": "Text-based", "pages": num_pages}
        else: # Very little text overall, default to image-based
            return {"type": "Image-based", "pages": num_pages}
        # --- End Classification Logic ---

    except Exception as e:
        print("Error: Processing PDF failed - {e}")
        if doc: # Ensure closure even if error happens during processing
            doc.close()
        return f"Error: Processing PDF failed - {e}"


# 定义POST请求路由
@app.post("/pdftypejudge/upload/")
async def upload_file(file: UploadFile = File(...)):
    """
    处理文件上传的POST请求
    """
    try:
        # 读取上传的文件内容
        contents = await file.read()
        # 确保目标目录存在
        save_dir = "/home/huangjunren/workflow-fastapi/test-data/"
        os.makedirs(save_dir, exist_ok=True)
        
        # 构建文件保存路径
        file_path = os.path.join(save_dir, file.filename)
        
        # 将文件内容写入目标路径
        with open(file_path, "wb") as f:
            f.write(contents)
            
        # 返回完整的文件路径供后续使用
        full_path = os.path.abspath(file_path)

        # 使用PyMuPDF判断PDF类型
        res = classify_pdf_pymupdf(full_path)
        print(f"PDF类型: {res}")

        # 创建PDF目标目录
        pdf_save_dir = "/home/huangjunren/workflow-fastapi/test-data/pdf/"
        os.makedirs(pdf_save_dir, exist_ok=True)
        
        # 构建新的文件路径
        # 生成随机6位英文+数字编码
        
        random_code = ''.join(random.choices(string.ascii_letters + string.digits, k=6))
        current_time = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        new_filename = f"{random_code}_{current_time}.pdf"
        
        new_file_path = os.path.join(pdf_save_dir, new_filename)
        
        # 移动文件到新路径
        os.rename(full_path, new_file_path)
        
        # 更新文件路径
        full_path = new_file_path

        # 返回文件信息
        # 将文件路径添加到返回结果中
        if isinstance(res, dict):
            res["file_path"] = full_path
            res["new_filename"] = new_filename
        else:
            res = {"msg": res}
        return res
    except Exception as e:
        return {"error": str(e)}

# 定义获取图片的GET请求路由
@app.get("/pdftypejudge/get_image/{image_name}")
async def get_image(image_name: str):
    try:
        # 构建图片完整路径
        image_path = os.path.join("/home/huangjunren/workflow-fastapi/test-data/images/", image_name)
        
        # 检查文件是否存在
        if not os.path.exists(image_path):
            return {"error": f"图片不存在: {image_name}"}
            
        # 检查文件是否为PNG格式
        if not image_name.lower().endswith('.png'):
            return {"error": "只支持PNG格式图片"}
            
        # 从文件系统读取图片
        return FileResponse(
            image_path,
            media_type="image/png",
            filename=image_name
        )
        
    except Exception as e:
        return {"error": str(e)}


# 定义简单的POST请求路由
@app.post("/pdftypejudge/transfrom_image/")
async def process_data(data: Dict):
    try:
        # 构建文件保存路径
        file_path = data.get("file_path")
        # 判断文件是否存在
        if not os.path.exists(file_path):
            return {"error": f"文件不存在: {file_path}"}

        wait_read_page_number = data.get("wait_read_page_number")
        if wait_read_page_number is None:
            return {"error": "未指定要读取的页码"}

        # 创建保存图片的目录
        image_save_dir = "/home/huangjunren/workflow-fastapi/test-data/images/"
        os.makedirs(image_save_dir, exist_ok=True)

        # 打开PDF文件
        doc = fitz.open(file_path)

        # 检查页码是否有效
        if wait_read_page_number < 0 or wait_read_page_number >= doc.page_count:
            doc.close()
            return {"error": f"页码 {wait_read_page_number} 超出范围 (0-{doc.page_count-1})"}

        # 获取指定页面
        page = doc.load_page(wait_read_page_number)
        
        # 将页面转换为高分辨率图片
        # 设置较高的分辨率参数，matrix用于缩放，这里设置为3倍清晰度
        zoom_x = 2.0  # 水平缩放
        zoom_y = 2.0  # 垂直缩放
        matrix = fitz.Matrix(zoom_x, zoom_y)
        pix = page.get_pixmap(matrix=matrix)
        
        # 构建输出图片路径
        output_filename = f"{os.path.splitext(os.path.basename(file_path))[0]}_page_{wait_read_page_number}.png"
        output_path = os.path.join(image_save_dir, output_filename)
        
        # 保存图片
        pix.save(output_path)
        
        # 关闭PDF文件
        doc.close()
        
        return {
            "success": True,
            "message": "PDF页面已成功转换为图片",
            "image_path": output_filename
        }
        
    except Exception as e:
        return {"error": str(e)}

# 启动服务器
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=6000)
